import sys
sys.path.append("..")

def discard_lines_before_headers(genofile):
    """Check for '[Data]' on line preceding expid line"""
    gf = open(genofile)
    linenum=0
    while True:
        line = gf.readline()
        linenum = linenum + 1
        if '[Data]' in line:
            break
        if not line or linenum > 10000:
            import sys
            print "Couldn't find string '[Data]' in the first %s lines of %s"%(linenum, genofile)
            print "Is it a correctly formatted Illumina genotype file?"
            sys.exit(1)
    gf.close()
    return linenum

def data_indices(genofile):
    gf = open(genofile)
    discard_lines = discard_lines_before_headers(genofile)
    for i in range(discard_lines):
        gf.readline()
    headers = gf.readline().split(',')
    for h, i in zip(headers, range(len(headers))):
        h = h.lstrip('"').rstrip('"')
        if "Sample ID" == h:
            sampleid = i
        if "Allele1 - Forward" == h:
            allele1 = i
        if "Allele2 - Forward" == h:
            allele2 = i
        if "SNP Name" == h:
            snp = i
    gf.close()
    return {'sampleid':sampleid, 'allele1':allele1, 'allele2':allele2, 'snp':snp}

def remove_duplicates(l):
    known = set()
    newlist = []

    for x in l:
      if x in known: continue
      newlist.append(x)
      known.add(x)

    l[:] = newlist

def return_sampleids(genofile, newgenofile):
    index_dict = data_indices(genofile)
    sampleid_index = index_dict['sampleid']
    gf = open(genofile)
    # discard lines before calls
    discard_lines = discard_lines_before_headers(genofile) + 1
    for i in range(discard_lines):
        gf.readline()
    i = 0
    idlist = []
    for line in gf:
        i = i+1
        line = line.lstrip().rstrip().split(',')
        sampleid = line[sampleid_index].lstrip('"').rstrip('"')
        idlist.append(sampleid)
    remove_duplicates(idlist)
    return idlist

def firstsampleid(genofile):
    index_dict = data_indices(genofile)
    sampleid_index = index_dict['sampleid']
    gf = open(genofile)
    # discard lines before calls
    discard_lines = discard_lines_before_headers(genofile) + 1
    for i in range(discard_lines):
        gf.readline()
    line = gf.readline().lstrip().rstrip().split(',')
    return line[sampleid_index].lstrip('"').rstrip('"')

def return_snps(genofile):
    index_dict = data_indices(genofile)
    sampleid_index = index_dict['sampleid']
    allele1_index = index_dict['allele1']
    allele2_index = index_dict['allele2']
    snp_index = index_dict['snp']
    gf = open(genofile)
    # discard lines before calls
    discard_lines = discard_lines_before_headers(genofile) + 1
    for i in range(discard_lines):
        gf.readline()
    fsi = firstsampleid(genofile)
    i = 0
    snplist = []
    while True:
        line = gf.readline()
        line = line.lstrip().rstrip().split(',')
        sampleid = line[sampleid_index].lstrip('"').rstrip('"')
        if sampleid != fsi:
            break
        snp =  line[snp_index].lstrip('"').rstrip('"')
        allele1 =  line[allele1_index].lstrip('"').rstrip('"')
        allele2 =  line[allele2_index].lstrip('"').rstrip('"')
        snplist.append(snp)
    return snplist

from utils import print_timing

@print_timing
def convert(genofile, newgenofile):
    index_dict = data_indices(genofile)
    sampleid_index = index_dict['sampleid']
    allele1_index = index_dict['allele1']
    allele2_index = index_dict['allele2']
    snp_index = index_dict['snp']
    sampleids = return_sampleids(genofile, newgenofile)
    snps = return_snps(genofile)
    snplen = len(snps)
    gf = open(genofile)
    ngf = open(newgenofile, 'w')
    # discard lines before calls
    discard_lines = discard_lines_before_headers(genofile)
    for i in range(discard_lines):
        line = gf.readline()
        ngf.write(line)
    header = gf.readline() # Read and discard headers
    sampleidsline = '\t' + '\t'.join(sampleids) + '\n'
    ngf.write(sampleidsline)
    # each consecutive block of calls corresponds to the same sampleid
    calls = []
    for line in gf:
        line = line.strip().split(',')
        sampleid = line[sampleid_index].lstrip('"').rstrip('"')
        allele1 = line[allele1_index].lstrip('"').rstrip('"')
        allele2 = line[allele2_index].lstrip('"').rstrip('"')
        snp     = line[snp_index].lstrip('"').rstrip('"')
        allele  = allele1 + allele2
        calls.append(allele)
    for snp, i in zip(snps, range(snplen)):
        if calls[i::snplen] != ['--']*len(sampleids):
            ngf.write('\t'.join([snp] + calls[i::snplen])+'\n')
    gf.close()
    ngf.close()

if __name__=='__main__':
    import sys
    from optparse import OptionParser
    usage = "usage is of the form '%prog originalgenofile convertedgenofile'"
    parser = OptionParser(usage=usage)
    (options, args) = parser.parse_args()

    if len(args) != 2:
        print "incorrect number of arguments"
        parser.print_help()
        sys.exit()

    genofile = args[0]
    convertedgenofile = args[1]
    convert(genofile, convertedgenofile)
